import csv
import numpy as np
import open3d as o3d


import os
import os.path as osp

RADIO = 80
MEASURE_UNIT = 'CM'


CSV_DIR = '../data/3d_csv'
TEST_CSV = './data/3d_csv/210423_123643_height.csv'
TEST_PCD = './data/3d_pcd/no_welding_40.pcd'
# TEST_CSV = './test.csv'



def get_csv_files(dir=CSV_DIR):
    return [osp.join(dir, file) for file in os.listdir(dir)]

def read_3d_csv(csv_file=TEST_CSV):
    x_points = []
    y_points = []
    z_points = []
    pcd = o3d.geometry.PointCloud()
    with open(csv_file) as f:
        reader = csv.reader(f)
        for col_index, row_obj in enumerate(reader):
            for row_index, item in enumerate(row_obj):
                if float(item) < 0:
                    item = 0
                    continue
                else:
                    x_points.append(row_index / RADIO)
                    y_points.append(col_index / RADIO)
                    z_points.append(item)
    # return x_points, y_points, z_points
    m3d = np.array([x_points, y_points, z_points]).T
    pcd.points = o3d.utility.Vector3dVector(m3d)
    return pcd

def vis_pcd(pcd, down_sample=True):
    if down_sample:
        voxel_down_pcd = pcd.voxel_down_sample(voxel_size=0.02)
        o3d.visualization.draw_geometries([voxel_down_pcd])
    else:
        o3d.visualization.draw_geometries([pcd])

def read_3d_csv_np(csv_file=TEST_CSV, is_filter=False):
    pcd = o3d.geometry.PointCloud()
    csv_data = np.loadtxt(csv_file, dtype=float, delimiter=',')
    print(csv_data.shape)
    y_num, x_num = csv_data.shape
    # 过滤-9999的点
    """
    考虑让<0的纵坐标变成 >0的较小的100个纵坐标的平均值
    """

    csv_data[csv_data < 0] = csv_data[csv_data> 0].min()
    # csv_data[csv_data < 0] = 0
    unit_x_point = np.arange(x_num)
    x_points = np.tile(unit_x_point, y_num)
    x_points = x_points/RADIO

    unit_y_point = np.arange(y_num)
    y_points = np.repeat(unit_y_point, x_num)/RADIO

    x_points, y_points, z_points = data_filter_csv(x_points, y_points, csv_data, is_filter)

    # add bottom
    # z_tem = z_points[z_points!=0]
    # z_bottom = np.full_like(y_points,z_tem.min()+0.42)
    # x_points = np.tile(x_points, 2).reshape(-1)
    # y_points = np.tile(y_points, 2).reshape(-1)
    # z_points = np.array([z_points,z_bottom]).reshape(-1)
    # z_points[z_points<z_tem.min()+0.42] = z_tem.min()+0.42

    m3d = np.array([x_points, y_points, z_points]).T
    pcd.points = o3d.utility.Vector3dVector(m3d)
    return pcd

def display_inlier_outlier(cloud, ind):
    """
    展示移除的点在原pcd的位置，显示为红色
    :param cloud: 点云pcd
    :param ind: 索引
    :return:
    """
    inlier_cloud = cloud.select_by_index(ind)
    outlier_cloud = cloud.select_by_index(ind, invert=True)

    print("Showing outliers (red) and inliers (gray): ")
    outlier_cloud.paint_uniform_color([1, 0, 0])
    inlier_cloud.paint_uniform_color([0.8, 0.8, 0.8])
    o3d.visualization.draw_geometries([inlier_cloud, outlier_cloud])

def read_pcd(pcd_file=TEST_PCD):
    pcd = o3d.io.read_point_cloud(pcd_file)
    return pcd

def data_filter_csv(x_points, y_points, z_points, is_filter=True):
    """get the in  interested area (x,y,z)

    :param x_points: [] (n,)
    :param y_points: [] (n,)
    :param z_points: [] (y_num, x_num)
    :param is_filter:
    :return:
        x_points[] (n,)
        y_points[] (n,)
        z_points[] (n,)
    """

    y_num, x_num = z_points.shape
    if is_filter:
        y_area_0 = y_num // 7
        x_area_0 = x_num // 3

        x_points = x_points.reshape(y_num, x_num)
        y_points = y_points.reshape(y_num, x_num)

        x_points = x_points[:y_area_0, :x_area_0]
        y_points = y_points[:y_area_0, :x_area_0]
        z_points = z_points[:y_area_0, :x_area_0]
        return x_points.reshape((-1)), y_points.reshape((-1)), z_points.reshape((-1))
    return x_points, y_points, z_points.reshape((-1))


def display_point_other(point):
    pass


def read_3d_npy(npy_file):
    """
    显示npy格式的三维点云数据
    :param npy_file: npy文件路径
    :return: o3d.geometry.PointCloud
    """

    z_points = np.load(npy_file)
    y_num, x_num = z_points.shape

    unit_x_point = np.arange(x_num)
    x_points = np.tile(unit_x_point, y_num) / RADIO

    unit_y_point = np.arange(y_num)
    y_points = np.repeat(unit_y_point, x_num) / RADIO

    z_points = z_points.reshape((-1))

    nan_index = np.isnan(z_points)

    # remove nan points
    x_points = x_points[~nan_index]
    y_points = y_points[~nan_index]
    z_points = z_points[~nan_index]

    pcd = o3d.geometry.PointCloud()
    m3d = np.array([x_points, y_points, z_points]).T
    pcd.points = o3d.utility.Vector3dVector(m3d)
    return pcd


def parse_npy(npy_file):
    """
    显示npy格式的三维点云数据
    :param npy_file: npy文件路径
    :return: (x_points->np.array, y_points, z_points) -> tuple
    """

    z_points = np.load(npy_file)
    y_num, x_num = z_points.shape

    unit_x_point = np.arange(x_num)
    x_points = np.tile(unit_x_point, y_num) / RADIO

    unit_y_point = np.arange(y_num)
    y_points = np.repeat(unit_y_point, x_num) / RADIO

    z_points = z_points.reshape((-1))

    nan_index = np.isnan(z_points)

    # remove nan points
    x_points = x_points[~nan_index]
    y_points = y_points[~nan_index]
    z_points = z_points[~nan_index]

    return (x_points, y_points, z_points)


if __name__ == '__main__':
    # pcd vis
    # pcd_file = './data/20210913_data/OK/OKfront/1.pcd'
    # pcd_file = '../data/20210913_data/NG/NGfront/1.pcd'
    # pcd = o3d.io.read_point_cloud(pcd_file)
    # vis_pcd(pcd)

    # csv vis
    # csv_file = './data/3d_csv/210423_160259_height.csv'
    # pcd = read_3d_csv(csv_file)
    # vis_pcd(pcd)

    # 3505277134ML085B085301119072200509_up_2022-09-15-08-48-25.npy
    npy_file = r'D:\hycx_work\point_cloud-deal\detection_3d\data\Side_Veri\3505277134ML085B085301129102100158_side_2021-12-14-16-27-31.npy'
    npy_file = 'D:\\35051772603505277134A16092200198251_up_2022-09-15-20-08-04.npy'
    npy_file ='F:\\big_tin\\2022-9-15-1\\3505277134ML085B085301119072200509_up_2022-09-15-08-48-25.npy'
    npy_file ='F:\\big_tin\\2022-9-15-1\\3505277134ML085B085301119072200407_up_2022-09-15-09-28-35.npy'
    npy_file ='F:\\big_tin\\2022-9-15-1\\3505277134ML085B085301119072200422_up_2022-09-15-05-46-45.npy'
    npy_file ='F:\\big_tin\\2022-9-15-1\\3505277134ML085B085301119072200423_up_2022-09-15-05-34-33.npy'
    npy_file ='F:\\big_tin\\2022-9-15-1\\3505277134ML085B085301119072200420_up_2022-09-15-06-38-12.npy'
    # npy_file = '../data/2021-12-14/3505277134ML085B085301129102100158_side_2021-12-14-17-42-57.npy'
    npy_file="D:\\Programs\\data\\dataset\\20220808+20220706_fly_tin\\generate_fly_tin\\45.npy"
    pcd = read_3d_npy(npy_file)
    vis_pcd(pcd)
